// MinNet.cpp : 此文件包含 "main" 函数。程序执行将在此处开始并结束。
//

#include "Util.hpp"
#include "Dataset.hpp"
#include "Function.hpp"
#include "Layer.hpp"
#include "Optimizer.hpp"

#include <iostream>
#include <string>
#include <time.h>

class RNN :public minnet::model::Model {
public:
    RNN() {
        lstm = minnet::model::LSTM(26, 26, 18, 1);
        register_layer(lstm);
    }
    minnet::Tensor Forward(minnet::Tensor& input) override {
        return lstm(input);
    }
private:
    minnet::model::LSTM lstm;
};

#define NUM 20074

int main(int argc, char** argv) {
    if(argc != 2){
        std::cout<<"must input name dataset path"<<std::endl;
        return 0;
    }
    
    std::string dataset_root_path = std::string(argv[1]);

    auto src_data = getAllName(dataset_root_path);
    if (src_data.size() != NUM) {
        std::cout<<"read dataset error!"<<std::endl;
        return 0;
    }
    srand((unsigned)time(NULL));
    std::vector<std::vector<float>> data(NUM);
    std::vector<std::vector<float>> label(NUM);

    for (int i = 0; i < NUM; i++) {
        data[i] = name_to_vec(src_data[i].second);
        label[i] = std::vector<float>(18, 0.f);
        label[i][src_data[i].first] = 1.f;
    }

    RNN net;
    float lr = 0.001f;
    minnet::optimizer::SGD opt(net.parameters(), lr, 0.9f);
    minnet::util::Timer timer;
    auto name_map = get_name_map();
    net.train();
    for (int i = 0; i < 10 ; i++) {
        minnet::shuffle(&data, &label, &src_data);
        float forward_time = 0.f;
        float backward_time = 0.f;
        float total_loss = 0.f;
        float temp_loss = 0.f;
        int count = 0;
        for (int j = 0; j < NUM; j++) {
            minnet::Tensor in;
            in.from_vector_1d(data[j], 0, -1);
            minnet::Tensor real;
            real.from_vector_1d(label[j], 0, -1);
            minnet::Tensor result = in.reshape(-1, 26);

            timer.begin();
            result = net(result);
            real = real.reshape(1, 18);
            if (minnet::argMax(result) == minnet::argMax(real)) count++;
            minnet::Tensor loss = minnet::function::CrossEntropyLoss(result, real);
            temp_loss += loss.at(0, 0);
            timer.end();
            forward_time += timer.cost();

            timer.begin();
            opt.zero_grad();
            loss.backward();
            opt.step();
            timer.end();
            backward_time += timer.cost();

            if ((j + 1) % 5000 == 0) {
                std::cout<< "epoch: "<< i + 1 <<" train(" << j + 1 << "/" << NUM << ")" << " loss: " << 
                    temp_loss / (5000.f) <<std::endl;
                total_loss += temp_loss;
                temp_loss = 0.f;
            }
        }
        if ((i + 1) % 3 == 0) opt.set_lr(opt.get_lr() / 2);
        std::cout << "epoch: " << i + 1 << " forward cost: " << forward_time << " backward cost: " << backward_time <<
            " avg loss: " << total_loss / (NUM * 1.f) << " accuracy in train set: " << count / (NUM * 1.f) << std::endl;
    }
    while (1) {
        std::cout << "Enter you name or Enter q quit:";
        std::string name;
        std::cin >> name;
        if (name == std::string("q")) break;
        std::cout << std::endl;
        auto name_vec = name_to_vec(name);
        minnet::Tensor in;
        in.from_vector_1d(name_vec, 0, -1);
        minnet::Tensor result = in.reshape(-1, 26);
        result = net(result);
        std::cout << "You must a " << (name_map.find(minnet::argMax(result)))->second << " people." << std::endl;
    }

    return 0;
}